You write custom CUDA kernels to replace pytorch operators in given architecture to get speedups. You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.

Here's an example to show you the syntax of inline embedding custom CUDA operators in PyTorch. The example given architecture is a simple addition:

python
import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
def init(self) -> None:
super().init()

def forward(self, a, b):
    return a + b
def get_inputs():
a = torch.randn(1, 128).cuda()
b = torch.randn(1, 128).cuda()
return [a, b]

def get_init_inputs():
return []



The example new architecture with a custom CUDA kernel looks like this:

python
import torch
from torch.utils.cpp_extension import load_inline

add_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

global void add_kernel(const float* a, const float* b, float* out, int size) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < size) {
out[idx] = a[idx] + b[idx];
}
}

torch::Tensor add_cuda(torch::Tensor a, torch::Tensor b) {
auto out = torch::empty_like(a);
int size = a.numel();
const int block_size = 256;
int num_blocks = (size + block_size - 1) / block_size;
add_kernel<<<num_blocks, block_size>>>(a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), size);
return out;
}
"""

add_cpp_source = """
torch::Tensor add_cuda(torch::Tensor a, torch::Tensor b);
"""

Compile the inline CUDA code
add = load_inline(
name="add",
cpp_sources=add_cpp_source,
cuda_sources=add_source,
functions=["add_cuda"],
verbose=True
)

class ModelNew(torch.nn.Module):
def init(self):
super(ModelNew, self).init()
self.add = add

def forward(self, a, b):
    return self.add.add_cuda(a, b)


---

Now, you are given the following PyTorch architecture to accelerate. The model computes the Dice Loss directly from 2D input tensors (e.g., from a convolutional layer) by first flattening them and then performing the loss calculation. This baseline implementation is efficient and uses PyTorch's highly optimized built-in functions for correctness and performance.

python
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    PyTorch基准实现：直接从2D输入计算Dice Loss
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """
        Applies Dice Loss to the 2D prediction and target tensors.

        Args:
            pred (torch.Tensor): Prediction tensor of shape [N, C, H, W].
            target (torch.Tensor): Target tensor of same shape as pred.

        Returns:
            torch.Tensor: Dice Loss value (scalar).
        """
        # --- 第一步：在内部展平张量 ---
        pred_flat = pred.view(-1)
        target_flat = target.view(-1)
        
        # --- 第二步：计算Dice Loss ---
        intersection = (pred_flat * target_flat).sum()
        pred_sum = pred_flat.sum()
        target_sum = target_flat.sum()
        
        epsilon = 1e-6
        dice_score = (2.0 * intersection) / (pred_sum + target_sum + epsilon)
        dice_loss = 1.0 - dice_score
        
        return dice_loss

batch_size = 32
height, width = 256, 256
channels = 1

def get_inputs():
    pred = torch.rand(batch_size, channels, height, width)
    target = torch.randint(0, 2, (batch_size, channels, height, width), dtype=torch.float32)
    return [pred, target]

def get_init_inputs():
    return []  # No special initialization inputs needed



Your task is to generate the `ModelNew` architecture with a custom CUDA kernel that fuses the `flatten` operation with the Dice Loss calculation. The implementation must be highly optimized.

**CRITICAL REQUIREMENTS:**

1.  **Performance Optimization & Fusion:**
    *   **Operator Fusion:** The entire calculation (flattening the input tensors and then computing the Dice Loss) must be performed within a **single CUDA kernel**. The kernel should take the multi-dimensional tensors as input and directly output the scalar loss value.
    *   The kernel should use a **Warp-level reduction** strategy for optimal performance. Warps should collaboratively iterate over the total number of elements in the tensors, effectively performing an "implicit flatten" within the kernel.

2.  **Kernel Logic:**
    *   The kernel should launch a grid of blocks where the total number of threads is sufficient to cover all elements in the input tensors (`pred.numel()`).
    *   Each thread should process one element of the flattened tensors.
    *   Use `__shfl_down_sync` for efficient warp-level reduction to compute the three required sums: `intersection`, `pred_sum`, and `target_sum`.
    *   After the kernel completes, the host-side C++ wrapper function should perform the final scalar arithmetic to compute the Dice Loss from the three sums.

3.  **Code Structure:** Follow the exact structure of the provided example, including `load_inline`, the CUDA source string, the C++ wrapper source string, and the `ModelNew` class. The `get_init_inputs` function must return `[]` to match the baseline.

4.  **Compilation Flags:** Use `-O3` and `--use_fast_math` for maximum performance, as this is a common practice for high-throughput kernels like this. Avoid hardcoding compute capabilities to ensure portability.
